package nl.utwente.eemcs.decomposition.transformation;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import nl.utwente.eemcs.graph.ActivityNode;
import nl.utwente.eemcs.graph.BranchedConstruct;
import nl.utwente.eemcs.graph.CommunicatorNode;
import nl.utwente.eemcs.graph.CommunicatorType;
import nl.utwente.eemcs.graph.ConditionalConstruct;
import nl.utwente.eemcs.graph.Edge;
import nl.utwente.eemcs.graph.EdgeType;
import nl.utwente.eemcs.graph.Graph;
import nl.utwente.eemcs.graph.GraphCollaboration;
import nl.utwente.eemcs.graph.LoopConstruct;
import nl.utwente.eemcs.graph.Node;
import nl.utwente.eemcs.graph.ParallelConstruct;
import nl.utwente.eemcs.graph.Partition;
import nl.utwente.eemcs.graph.distribution.DistributionList;
import nl.utwente.eemcs.graph.distribution.DistributionLocation;
import nl.utwente.eemcs.graphviz.util.Debugger;

/**
 * Decomposer class decomposes a graph based upon a distribution list into
 * separate processes for distribution on-premise or in the cloud The result of
 * the transform function is a collaborative model with separate processes and
 * communication edges between the processes. NOTE: The transformation process
 * performs the transformations by changing the source graph, instead of copying
 * the original graph.
 * 
 * @author Evert Duipmans
 */
public class Decomposer {
   /**
    * Unique counter, used during creation of new elements
    */
   private int uniqueCounter = 0;

   /**
    * Main graph on which the transformations are performed
    */
   private Graph inputGraph;

   /**
    * List with all the separate process graphs (including the main graph and
    * the branches of composite constructs)
    */
   private List<Graph> graphs = new ArrayList<Graph>();

   /**
    * List with all the composite constructs
    */
   private List<BranchedConstruct<?, ?>> constructs = new ArrayList<BranchedConstruct<?, ?>>();

   /**
    * List with all the loop constructs
    */
   private List<LoopConstruct> loopConstructs = new ArrayList<LoopConstruct>();

   /**
    * Distribution list with distribution locations for activities
    */
   private DistributionList distributionList;

   /**
    * Initialize a decomposer object with a graph and a distribution list
    * 
    * @param input
    *           Graph which should be transformed
    * @param distributionList
    *           List with distribution locations and data dependencies
    */
   private Decomposer(Graph input, DistributionList distributionList) {
      this.inputGraph = input;
      this.distributionList = distributionList;
   }

   /**
    * Performs all the necessary steps to transform the process NOTE: The
    * transformation process performs the transformations by changing the source
    * graph, instead of copying the original graph.
    */
   public static GraphCollaboration performDecomposition(Graph graph,
         DistributionList distributionList) throws DecompositionException {
      Decomposer decomposer = new Decomposer(graph, distributionList);

      // Step 1: walk through the process and mark activities with their
      // distribution location and collect the composite constructs
      decomposer.collectSubProcesses();
      // Step 2: partition adjacent activities with the same distribution
      // location
      decomposer.partitionGraphs();
      // Step 3: create communicators between partitions and merge partitions
      // decomposer.createCommunicators();
      decomposer.createCommunicators2();
      // Step 4: create collaborative model with processes and communication
      // edges
      return decomposer.collectOutputProcesses();
   }

   /**
    * Marks the original graph with the distribution locations and collects the
    * branches of the parallel and conditional constructs
    */
   public void collectSubProcesses() throws DecompositionException {
      // When there is a start node in the graph, start parsing
      if (inputGraph.getStartNode() != null) {
         uniqueCounter = 0;
         graphs.clear();
         graphs.add(inputGraph);
         collectSubProcesses(inputGraph.getStartNode(), inputGraph);

         Debugger.drawImage(graphs, "decomposition_collectSubProcesses");
      } else
         throw new DecompositionException("The input graph has no start node");
   }

   /**
    * Walks through the graph and marks the nodes with its distribution location
    * and collects all the subgraphs (branches)
    * 
    * @param node
    *           The current node that is processed
    * @param newGraph
    *           The newly created graph
    * @return
    */
   private void collectSubProcesses(Node node, Graph originalGraph)
         throws DecompositionException {
      if (node instanceof BranchedConstruct<?, ?>) {
         constructs.add((BranchedConstruct<?, ?>) node);

         // Set the distribution location of the nodes
         Node startNode = ((BranchedConstruct<?, ?>) node).getStartNode();
         startNode.setLocation(distributionList
               .getDistributionLocation(startNode));
         ((BranchedConstruct<?, ?>) node).getEndNode().setLocation(
               distributionList.getDistributionLocation(startNode));
         node.setLocation(distributionList.getDistributionLocation(startNode));

         for (Graph branch : ((BranchedConstruct<?, ?>) node).getBranches()) {
            // When a branch has no nodes, skip to the next branch
            if (branch == null || branch.getNodes().size() == 0)
               continue;

            // Create temporary node
            workOnBranch(branch, node.getLocation());
         }
      } else if (node instanceof LoopConstruct) {
         LoopConstruct castedNode = (LoopConstruct) node;
         loopConstructs.add(castedNode);

         // Set distribution location
         castedNode.getConditionNode().setLocation(
               distributionList.getDistributionLocation(castedNode
                     .getConditionNode()));
         castedNode.setLocation(castedNode.getConditionNode().getLocation());

         // Create temporary node
         workOnBranch(castedNode.getInnerGraph(), node.getLocation());
      } else {
         // Set distribution location of the node
         if (distributionList.hasDistributionLocation(node))
            node.setLocation(distributionList.getDistributionLocation(node));
      }

      // Get outgoing edges and follow these edges
      List<Edge> outgoingEdges = originalGraph.getOutgoingEdges(node,
            EdgeType.Control,EdgeType.Communication);

      
      for(Edge e : outgoingEdges)
         collectSubProcesses(e.getEnd(), originalGraph);
      
//      if (outgoingEdges.size() == 1)
//         collectSubProcesses(outgoingEdges.get(0).getEnd(), originalGraph);
//      else if (outgoingEdges.size() > 1)
//         throw new DecompositionException("Node \"" + node.getName()
//               + "\" has multiple outgoing control edges, which is not allowed");
   }

   private void workOnBranch(Graph branch, DistributionLocation distLoc)
         throws DecompositionException {
      // Add branch to graph list
      this.graphs.add(branch);
      Node oldStartNode = branch.getStartNode();

      // Add new temporary node to the beginning of the branch, to mark the
      // distribution location of the composite construct
      ActivityNode newNode = new ActivityNode("temp" + uniqueCounter++);
      newNode.setLocation(distLoc);
      branch.addNode(newNode);
      branch.setStartNode(newNode);
      branch.addEdge(new Edge(newNode, oldStartNode, EdgeType.Control));

      // Parse branch
      collectSubProcesses(oldStartNode, branch);
   }

   /**
    * Partition adjacent node with the same distribution location. A partition
    * is a node with an internal graph
    */
   public void partitionGraphs() {
      uniqueCounter = 0;
      // Walk through each of the graphs
      for (Graph g : this.graphs) {
         Node startNode = g.getStartNode();

         // Create first partition
         Partition p = new Partition("partition" + uniqueCounter++);
         p.setLocation(startNode.getLocation());
         g.addNode(p);
         g.setStartNode(p);

         // Partition the graph
         partitionGraph(startNode, g, p);
      }

      Debugger.drawImage(graphs, "decomposition_partitioned");
   }

   /**
    * Partition adjacent nodes with the same distribution location
    * 
    * @param node
    *           Current node
    * @param g
    *           Parent graph of the node
    * @param currentPartition
    *           The current partition in which the node will be assigned
    */
   private void partitionGraph(Node node, Graph g, Partition currentPartition) {
      // Add node to partition
      currentPartition.addNode(node);

      // Remove node from graph
      g.removeNode(node);

      // Get outgoing edges and follow these edges
      List<Edge> outgoingEdges = g.getOutgoingEdges(node, EdgeType.Control, EdgeType.Communication);

      //if (outgoingEdges.size() == 1) {
        // Edge e = outgoingEdges.get(0);
      for(Edge e : outgoingEdges) {
         // Remove edge from graph
         g.removeEdge(e);

         // Compare distribution location of the current node and the next node
         if (e.getEnd().getLocation() == node.getLocation()) {
            // The same distribution location:
            // Add edge to current partition, continue walking and use the same
            // partition
            currentPartition.addEdge(e);
            partitionGraph(e.getEnd(), g, currentPartition);
         } else {
            // Different distribution location:
            // Create new partition, create edge between partitions and add to
            // parent graph
            Partition newPartition = new Partition("partition"
                  + uniqueCounter++);
            newPartition.setLocation(e.getEnd().getLocation());
            g.addNode(newPartition);
            g.addEdge(new Edge(currentPartition, newPartition, EdgeType.Control));

            // Continue and use new partition
            partitionGraph(e.getEnd(), g, newPartition);
         }
      }
   }

   public void createCommunicators2() {
      // Reset the unique counter
      uniqueCounter = 0;
      // For each graph,
      for (Graph g : this.graphs) {
         createCommunicators2(g);

         // Increase unique counter
         // uniqueCounter++;
      }

      Debugger.drawImage(graphs, "decomposition_communicators");
   }

   private void createCommunicators2(Graph currentGraph) {
      if (currentGraph.getNodes().size() == 0)
         return;

      // Create lists for odd and even partitions
      List<Partition> oddPartitions = new ArrayList<Partition>();
      List<Partition> evenPartitions = new ArrayList<Partition>();

      // Add start partition
      oddPartitions.add((Partition) currentGraph.getStartNode());

      // Collect even and odd partitions
      List<Edge> outgoingEdges = currentGraph.getOutgoingEdges(
            currentGraph.getStartNode(), EdgeType.Control);
      for (int i = 0; outgoingEdges.size() > 0; i++) {
         Partition nextPart = (Partition) outgoingEdges.get(0).getEnd();
         // Add to the correct list
         ((i % 2 == 0) ? evenPartitions : oddPartitions).add(nextPart);
         outgoingEdges = currentGraph.getOutgoingEdges(nextPart,
               EdgeType.Control);
      }

      if (evenPartitions.size() == 0)
         return;

      for (int i = 0; i < oddPartitions.size(); i++) {
         if (evenPartitions.size() <= i)
            continue;

         // Get partitions
         Partition oddPartition = oddPartitions.get(i);
         Partition evenPartition = evenPartitions.get(i);

         // Create communicator nodes
         CommunicatorNode invrecNode = new CommunicatorNode("InvokeReceive"
               + uniqueCounter, CommunicatorType.InvokeRec,
               oddPartition.getLocation());
         CommunicatorNode invresNode = new CommunicatorNode("InvokeResponse"
               + uniqueCounter, CommunicatorType.InvokeRes,
               oddPartition.getLocation());
         CommunicatorNode recNode = new CommunicatorNode("Receive"
               + uniqueCounter, CommunicatorType.Receive,
               evenPartition.getLocation());
         CommunicatorNode repNode = new CommunicatorNode("Response"
               + uniqueCounter++, CommunicatorType.Response,
               evenPartition.getLocation());

         // Create and add edges
         oddPartition.addEdge(new Edge(oddPartition.getEndNode(), invrecNode,
               EdgeType.Control));
         oddPartition
               .addEdge(new Edge(invrecNode, invresNode, EdgeType.Control));
         evenPartition.addEdge(new Edge(recNode, evenPartition.getStartNode(),
               EdgeType.Control));
         evenPartition.addEdge(new Edge(evenPartition.getEndNode(), repNode,
               EdgeType.Control));

         // Create and add communication edges
         currentGraph.addEdge(new Edge(invrecNode, recNode,
               EdgeType.Communication));
         currentGraph.addEdge(new Edge(repNode, invresNode,
               EdgeType.Communication));

         // Add nodes
         oddPartition.addNode(invrecNode);
         oddPartition.addNode(invresNode);
         evenPartition.addNode(recNode);
         evenPartition.getInnerGraph().setStartNode(recNode);
         evenPartition.addNode(repNode);
      }

      Partition oddPartition = oddPartitions.get(0);
      Partition evenPartition = evenPartitions.get(0);

      // Merge partitions
      for (int i = 1; i < oddPartitions.size(); i++) {
         Partition nextPart = oddPartitions.get(i);
         // Create connection edge
         oddPartition.getInnerGraph().addEdge(
               new Edge(oddPartition.getEndNode(), nextPart.getStartNode(),
                     EdgeType.Control));
         for (Node n : nextPart.getInnerGraph().getNodes())
            oddPartition.addNode(n);
         for (Edge e : nextPart.getInnerGraph().getEdges())
            oddPartition.addEdge(e);
      }

      for (int i = 1; i < evenPartitions.size(); i++) {
         Partition nextPart = evenPartitions.get(i);
         // Create connection edge
         evenPartition.getInnerGraph().addEdge(
               new Edge(evenPartition.getInnerGraph().getEndNode(), nextPart
                     .getStartNode(), EdgeType.Control));
         for (Node n : nextPart.getInnerGraph().getNodes())
            evenPartition.addNode(n);
         for (Edge e : nextPart.getInnerGraph().getEdges())
            evenPartition.addEdge(e);
      }

      // Remove control edges between partitions
      for (Edge e : currentGraph.getEdges(EdgeType.Control))
         currentGraph.removeEdge(e);
      currentGraph.getNodes().clear();

      currentGraph.addNode(oddPartition);
      currentGraph.addNode(evenPartition);
   }

   /**
    * Create communicators between processes
    */
   public void createCommunicators() {
      // Reset the unique counter
      uniqueCounter = 0;
      // For each graph,
      for (Graph g : this.graphs) {
         Partition startPartition = (Partition) g.getStartNode();
         createCommunicators(startPartition, g);

         // Increase unique counter
         uniqueCounter++;
      }

      Debugger.drawImage(graphs, "decomposition_communicators");
   }

   /**
    * Take first partition and create communicators from/to the second
    * partition. Merge the third partition and start the operation again on
    * partition 1.
    * 
    * @param currentPartition
    *           Current partition
    * @param currentGraph
    *           Graph in which the partition is placed
    */
   private void createCommunicators(Partition currentPartition,
         Graph currentGraph) {
      // If there is a next partition, create a invoke/receive communication
      // pair
      if (currentGraph.getOutgoingEdges(currentPartition, EdgeType.Control)
            .size() == 1) {
         Edge part1part2Edge = currentGraph.getOutgoingEdges(currentPartition,
               EdgeType.Control).get(0);
         Partition partition2 = (Partition) part1part2Edge.getEnd();

         // Remove old edge from p1 to p2
         currentGraph.removeEdge(part1part2Edge);

         /**
          * Create invRec/receive communicators
          */
         // Add nodes and edges to p1
         CommunicatorNode invrecNode = new CommunicatorNode("InvokeReceive"
               + uniqueCounter, CommunicatorType.InvokeRec,
               currentPartition.getLocation());
         currentPartition.addEdge(new Edge(currentPartition.getEndNode(),
               invrecNode, EdgeType.Control));
         currentPartition.addNode(invrecNode);

         // Add nodes and edges to p2
         CommunicatorNode recNode = new CommunicatorNode("Receive"
               + uniqueCounter, CommunicatorType.Receive,
               partition2.getLocation());
         partition2.addEdge(new Edge(recNode, partition2.getStartNode(),
               EdgeType.Control));
         partition2.addNode(recNode);
         partition2.getInnerGraph().setStartNode(recNode);

         // Add communication edges
         currentGraph.addEdge(new Edge(invrecNode, recNode,
               EdgeType.Communication));

         /**
          * Create response/invRep communicators
          */
         // Add nodes and edges to p2
         CommunicatorNode resNode = new CommunicatorNode("Response"
               + uniqueCounter, CommunicatorType.Response,
               partition2.getLocation());
         partition2.addEdge(new Edge(partition2.getEndNode(), resNode,
               EdgeType.Control));
         partition2.addNode(resNode);

         // Add nodes and edges to p1
         CommunicatorNode invresNode = new CommunicatorNode("InvokeResponse"
               + uniqueCounter, CommunicatorType.InvokeRes,
               currentPartition.getLocation());
         currentPartition.addNode(invresNode);
         currentPartition.addEdge(new Edge(invrecNode, invresNode,
               EdgeType.Control));

         // Add communication edges
         currentGraph.addEdge(new Edge(resNode, invresNode,
               EdgeType.Communication));

         /**
          * Combine partition 1 and 3 (if available)
          */
         if (currentGraph.getOutgoingEdges(partition2, EdgeType.Control).size() == 1) {
            Edge part2part3Edge = currentGraph.getOutgoingEdges(partition2,
                  EdgeType.Control).get(0);
            Partition partition3 = (Partition) part2part3Edge.getEnd();

            // Connect partition 1 with partition 3
            currentPartition.addEdge(new Edge(invresNode, partition3
                  .getStartNode(), EdgeType.Control));

            // Copy nodes and edges from p3 to p1
            for (Node n : partition3.getInnerGraph().getNodes())
               currentPartition.addNode(n);
            for (Edge e : partition3.getInnerGraph().getEdges())
               currentPartition.addEdge(e);

            // Update outgoing edges from partition 3
            for (Edge outgoingEdge : currentGraph.getOutgoingEdges(partition3,
                  EdgeType.Control))
               outgoingEdge.setStart(currentPartition);

            // Remove old edge from p2 to p3 and partition 3
            currentGraph.removeEdge(part2part3Edge);
            currentGraph.removeNode(partition3);

            // Increase counter and continue replacing
            uniqueCounter++;
            createCommunicators(currentPartition, currentGraph);
         }
      }
   }

   /**
    * Collect the separate processes and the communication edges between these
    * processes
    * 
    * @return Collaborative graph model with separate processes and
    *         communication edges
    */
   public GraphCollaboration collectOutputProcesses() {
      ArrayList<Graph> outputProcesses = new ArrayList<Graph>();
      Set<Edge> communicationEdges = new HashSet<Edge>();
      Set<Edge> dataEdges = new HashSet<Edge>();

      // Merge branches of composite constructs
      replaceCompositeNodeProcesses(outputProcesses, communicationEdges);

      // Collect processes
      int i = 0;
      for (Node n : inputGraph.getNodes())
         outputProcesses.add(i++, ((Partition) n).getInnerGraph());
      for (Edge e : inputGraph.getEdges()) {
         if (e.getEdgeType() == EdgeType.Communication)
            communicationEdges.add(e);
         else
            dataEdges.add(e);
      }

      return new GraphCollaboration(outputProcesses, communicationEdges,
            dataEdges, inputGraph);
   }

   /**
    * Merges branches of composite constructs, since the distribution location
    * of the partitions of the branches are always the same
    * 
    * @param outputProcesses
    *           List with output processes
    * @param communicationEdges
    *           List with communication edges
    */
   private void replaceCompositeNodeProcesses(ArrayList<Graph> outputProcesses,
         Set<Edge> communicationEdges) {
      for (BranchedConstruct<?, ?> p : this.constructs) {
         if (p instanceof ParallelConstruct)
            replaceCompositeNodeProcessesParallel((ParallelConstruct) p,
                  outputProcesses, communicationEdges);
         else if (p instanceof ConditionalConstruct)
            replaceCompositeNodeProcessesConditional((ConditionalConstruct) p,
                  outputProcesses, communicationEdges);
      }
      for (LoopConstruct c : this.loopConstructs)
         replaceLoopNodeProcess(c, outputProcesses, communicationEdges);
   }

   private void replaceCompositeNodeProcessesParallel(
         ParallelConstruct construct, ArrayList<Graph> outputProcesses,
         Set<Edge> communicationEdges) {
      List<Graph> branches = new ArrayList<Graph>();

      for (Graph g : construct.getBranches())
         branches.add(dealWithBranch(g, outputProcesses, communicationEdges));

      // Replace branches
      construct.setBranches(branches);
   }

   private void replaceCompositeNodeProcessesConditional(
         ConditionalConstruct construct, ArrayList<Graph> outputProcesses,
         Set<Edge> communicationEdges) {
      Graph trueBranch = null;
      Graph falseBranch = null;

      for (Graph g : construct.getBranches()) {
         // When the branch has no nodes, continue to the next
         if (g == null || g.getNodes().size() == 0)
            continue;

         if (construct.getBranch(true) == g)
            trueBranch = dealWithBranch(g, outputProcesses, communicationEdges);
         else
            falseBranch = dealWithBranch(g, outputProcesses, communicationEdges);
      }

      // Replace branches
      construct.setBranch(trueBranch, true);
      construct.setBranch(falseBranch, false);
   }

   private void replaceLoopNodeProcess(LoopConstruct construct,
         ArrayList<Graph> outputProcesses, Set<Edge> communicationEdges) {
      construct.setInnerGraph(dealWithBranch(construct.getInnerGraph(),
            outputProcesses, communicationEdges));
   }

   private Graph dealWithBranch(Graph branch, ArrayList<Graph> outputProcesses,
         Set<Edge> communicationEdges) {
      // Remove from graphs list
      graphs.remove(branch);

      // Set the branch of the construct
      Partition firstPartition = (Partition) branch.getStartNode();

      // Remove startNode and start edge
      Edge oldEdge = firstPartition
            .getInnerGraph()
            .getOutgoingEdges(firstPartition.getInnerGraph().getStartNode(),
                  EdgeType.Control).get(0);
      firstPartition.getInnerGraph().removeNode(
            firstPartition.getInnerGraph().getStartNode());

      // Get the control edge from the temporary start node and remove it from
      // the graph, reset the start node
      firstPartition.getInnerGraph().removeEdge(oldEdge);
      firstPartition.getInnerGraph().setStartNode(oldEdge.getEnd());

      for (Edge e : branch.getEdges())
         communicationEdges.add(e);

      // Copy partitions
      for (Node n : branch.getNodes()) {
         if (n != branch.getStartNode())
            if (n instanceof Partition)
               outputProcesses.add(((Partition) n).getInnerGraph());
      }

      return firstPartition.getInnerGraph();
   }
}